"""
使用gdal.warp实现hdf子数据集转tif
优化版本2.0
提高代码可读性和可维护性
保留原始数据集的投影信息
"""
import os
from osgeo import gdal, gdalconst

# 写入tif
def write_tif(file_path, data, geotransform, projInfo, nodata, gdal_type):
    # 相当于一个创建数据集的驱动
    driver = gdal.GetDriverByName("GTiff")
    # 根据数据的维度来确定行列数（图像大小）
    rows, cols = data.shape
    # 创建一个新的数据集，存储输出文件；1是波段
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    # 设置仿射变换矩阵
    dataset.SetGeoTransform(geotransform)
    # 设置投影
    dataset.SetProjection(projInfo)
    # 1个波段
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    # 释放内存
    del dataset

# 读取hdf文件并重投影为WGS84坐标系
# 输入 hdf_file: hdf文件路径
# 输出 tif_dir: tif文件保存目录
def hdf2tif(hdf_path, tif_dir):
    dataset_list = []   # 存放子数据集
    # EVI数据
    if 'MOD13A2' in hdf_path:
        # 需要获取的子数据集名称，和dataset_list一一对应
        subDatasetList = ['EVI', 'PR', 'QC']
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days EVI"'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days pixel reliability"'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days VI Quality"'.format(hdf_path)))
        for i, dataset in enumerate(dataset_list):
            nodata = dataset.GetRasterBand(1).GetNoDataValue()
            data = dataset.GetRasterBand(1).ReadAsArray()
            geotransform = dataset.GetGeoTransform()
            projInfo = dataset.GetProjection()
            if not os.path.exists(tif_dir + "/" + subDatasetList[i]):
                os.mkdir(tif_dir + "/" + subDatasetList[i])
            output_path = tif_dir + "/" + subDatasetList[i] + "/" + os.path.splitext(hdf_path.split("\\")[-1])[0] + "_{}.tif".format(subDatasetList[i])
            # gdal.Warp(output_path, dataset, dstSRS='EPSG:4326', srcNodata=nodata, dstNodata=nodata, resampleAlg=gdalconst.GRA_NearestNeighbour)
            if subDatasetList[i] == 'EVI':
                write_tif(output_path, data,  geotransform, projInfo, nodata, gdalconst.GDT_Int16)
            elif subDatasetList[i] == 'PR':
                write_tif(output_path, data,  geotransform, projInfo, nodata, gdalconst.GDT_UInt16)
            elif subDatasetList[i] == 'QC':
                write_tif(output_path, data,  geotransform, projInfo, nodata, gdalconst.GDT_Int16)
    # LST数据
    elif 'MOD11A1' in hdf_path:
        subDatasetList = ['LST_day', 'LST_day_qc', 'LST_night', 'LST_night_qc']
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_Daily_1km_LST:LST_Day_1km'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_Daily_1km_LST:QC_Day'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_Daily_1km_LST:LST_Night_1km'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_Daily_1km_LST:QC_Night'.format(hdf_path)))
        for i, dataset in enumerate(dataset_list):
            nodata = dataset.GetRasterBand(1).GetNoDataValue()
            data = dataset.GetRasterBand(1).ReadAsArray()
            geotransform = dataset.GetGeoTransform()
            projInfo = dataset.GetProjection()
            if not os.path.exists(tif_dir + "/" + subDatasetList[i]):
                os.mkdir(tif_dir + "/" + subDatasetList[i])
            output_path = tif_dir + "/" + subDatasetList[i] + "/" + os.path.splitext(hdf_path.split("\\")[-1])[0] + "_{}.tif".format(subDatasetList[i])
            # gdal.Warp(output_path, dataset, dstSRS='EPSG:4326', srcNodata=nodata, dstNodata=nodata, resampleAlg=gdalconst.GRA_NearestNeighbour)
            write_tif(output_path, data, geotransform, projInfo, nodata, gdalconst.GDT_UInt16)
    # 反照率数据
    elif 'GLASS02A06' in hdf_path:
        subDatasetList = ['ABD_BSA_VIS', 'ABD_WSA_VIS', 'QC_VIS']
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":GLASS02A06:ABD_BSA_VIS'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":GLASS02A06:ABD_WSA_VIS'.format(hdf_path)))
        dataset_list.append(gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":GLASS02A06:QC_VIS'.format(hdf_path)))
        for i, dataset in enumerate(dataset_list):
            nodata = dataset.GetRasterBand(1).GetNoDataValue()
            data = dataset.GetRasterBand(1).ReadAsArray()
            geotransform = dataset.GetGeoTransform()
            projInfo = dataset.GetProjection()
            if not os.path.exists(tif_dir + "/" + subDatasetList[i]):
                os.mkdir(tif_dir + "/" + subDatasetList[i])
            output_path = tif_dir + "/" + subDatasetList[i] + "/" + os.path.splitext(hdf_path.split("\\")[-1])[0] + "_{}.tif".format(subDatasetList[i])
            # gdal.Warp(output_path, dataset, dstSRS='EPSG:4326', srcNodata=nodata, dstNodata=nodata, resampleAlg=gdalconst.GRIORA_NearestNeighbour)
            if subDatasetList[i] == 'ABD_BSA_VIS' or subDatasetList[i] == 'ABD_WSA_VIS':
                write_tif(output_path, data, geotransform, projInfo, nodata, gdalconst.GDT_Int16)
            else:
                write_tif(output_path, data, geotransform, projInfo, nodata, gdalconst.GDT_UInt16)


def main(hdf_dir, tif_dir):
    basename = ".hdf"
    hdf_ls = [os.path.join(hdf_dir, f) for f in os.listdir(hdf_dir) if f.endswith(basename)]
    for hdf_path in hdf_ls:
        # 按照区域划分输出文件夹
        new_tif_dir = ''
        if 'h25v04' in hdf_path:
            new_tif_dir = tif_dir + "/h25v04"
        elif 'h25v05' in hdf_path:
            new_tif_dir = tif_dir + "/h25v05"
        if not os.path.exists(new_tif_dir):
            os.makedirs(new_tif_dir)
        hdf2tif(hdf_path, new_tif_dir)
        print("{} done".format(hdf_path))


if __name__ == '__main__':
    hdf_dir = r"G:\test\LST_hdf"
    tif_dir = r"G:\test\process_result\LST_tif_new"
    main(hdf_dir, tif_dir)
